{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sympy import symbols, cos, sin, pi, simplify, pprint, tan, expand_trig, sqrt, trigsimp, atan2\n",
    "from sympy.matrices import Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def pose(theta, alpha, a, d):\n",
    "  # returns the pose T of one joint frame i with respect to the previous joint frame (i - 1)\n",
    "  # given the parameters:\n",
    "  # theta: theta[i]\n",
    "  # alpha: alpha[i-1]\n",
    "  # a: a[i-1]\n",
    "  # d: d[i]\n",
    "\n",
    "  r11, r12 = cos(theta), -sin(theta)\n",
    "  r23, r33 = -sin(alpha), cos(alpha)\n",
    "  r21 = sin(theta) * cos(alpha)\n",
    "  r22 = cos(theta) * cos(alpha)\n",
    "  r31 = sin(theta) * sin(alpha)\n",
    "  r32 = cos(theta) * sin(alpha)\n",
    "  y = -d * sin(alpha)\n",
    "  z = d * cos(alpha)\n",
    "    \n",
    "  T = Matrix([\n",
    "    [r11, r12, 0.0, a],\n",
    "    [r21, r22, r23, y],\n",
    "    [r31, r32, r33, z],\n",
    "    [0.0, 0.0, 0.0, 1]\n",
    "  ])\n",
    "  \n",
    "  T = simplify(T)\n",
    "\n",
    "  return T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# get the pose (homogenous transforms) of each joint wrt to previous joint\n",
    "\n",
    "q1, q2, q3, q4, q5, q6= symbols('q1:7')\n",
    "\n",
    "d90 = pi / 2 \n",
    "\n",
    "T01 = pose(q1, 0, 0, 0.75)\n",
    "T12 = pose(q2 - d90, -d90, 0.35, 0)\n",
    "T23 = pose(q3, 0, 1.25, 0)\n",
    "T34 = pose(q4, -d90, -0.054, 1.5)\n",
    "T45 = pose(q5, d90, 0, 0)\n",
    "T56 = pose(q6, -d90, 0, 0)\n",
    "T6g = pose(0, 0, 0, 0.303)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R03 = \n",
      "\n",
      "Matrix([\n",
      "[sin(q2 + q3)*cos(q1), cos(q1)*cos(q2 + q3), -sin(q1)],\n",
      "[sin(q1)*sin(q2 + q3), sin(q1)*cos(q2 + q3),  cos(q1)],\n",
      "[        cos(q2 + q3),        -sin(q2 + q3),        0]])\n",
      "R03.T = \n",
      "\n",
      "Matrix([\n",
      "[sin(q2 + q3)*cos(q1), sin(q1)*sin(q2 + q3),  cos(q2 + q3)],\n",
      "[cos(q1)*cos(q2 + q3), sin(q1)*cos(q2 + q3), -sin(q2 + q3)],\n",
      "[            -sin(q1),              cos(q1),             0]])\n"
     ]
    }
   ],
   "source": [
    "# From the poses, get the rotation of joint 3 wrt to the base frame and the transpose \n",
    "# We will need this later\n",
    "\n",
    "\n",
    "T03 = simplify(T01 * T12 * T23)\n",
    "R03 = T03[:3, :3]\n",
    "R03T = R03.T\n",
    "\n",
    "print(\"R03 = \")\n",
    "print()\n",
    "print(R03)\n",
    "\n",
    "print(\"R03.T = \")\n",
    "print()\n",
    "print(R03T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R36 = \n",
      "\n",
      "Matrix([\n",
      "[-sin(q4)*sin(q6) + cos(q4)*cos(q5)*cos(q6), -sin(q4)*cos(q6) - sin(q6)*cos(q4)*cos(q5), -sin(q5)*cos(q4)],\n",
      "[                           sin(q5)*cos(q6),                           -sin(q5)*sin(q6),          cos(q5)],\n",
      "[-sin(q4)*cos(q5)*cos(q6) - sin(q6)*cos(q4),  sin(q4)*sin(q6)*cos(q5) - cos(q4)*cos(q6),  sin(q4)*sin(q5)]])\n"
     ]
    }
   ],
   "source": [
    "# From the poses, get the rotation of joint 6 wrt to the joint 3\n",
    "# We will need this later \n",
    "\n",
    "T36 = simplify(T34 * T45 * T56)\n",
    "R36 = T36[:3, :3]\n",
    "\n",
    "print(\"R36 = \")\n",
    "print()\n",
    "print(R36)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# rotation matrices in x, y, z axes\n",
    "\n",
    "def rotx(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [1., 0., 0.],\n",
    "    [0., cq,-sq],\n",
    "    [0., sq, cq]\n",
    "  ])\n",
    "    \n",
    "  return r\n",
    "\n",
    "\n",
    "def roty(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [ cq, 0., sq],\n",
    "    [ 0., 1., 0.],\n",
    "    [-sq, 0., cq]\n",
    "  ])\n",
    "    \n",
    "  return r\n",
    "\n",
    "\n",
    "def rotz(q):\n",
    "\n",
    "  sq, cq = sin(q), cos(q)\n",
    "\n",
    "  r = Matrix([\n",
    "    [cq,-sq, 0.],\n",
    "    [sq, cq, 0.],\n",
    "    [0., 0., 1.]\n",
    "  ])\n",
    "    \n",
    "  return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡ 0    0    1⎤\n",
      "⎢            ⎥\n",
      "⎢ 0   -1.0  0⎥\n",
      "⎢            ⎥\n",
      "⎣1.0   0    0⎦\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "# the yaw, pitch roll is given wrt to the URDF frame \n",
    "# We must convert this to gripper frame by performing\n",
    "# a rotation of 180 degrees ccw about the z axis and then \n",
    "# a rotation of 90 degrees cw about the new y axis\n",
    "\n",
    "# This is the transpose of the rotation of the urdf frame wrt to gripper frame and its transpose\n",
    "# ( which is strangely the same) which is important later\n",
    "\n",
    "Rgu = (rotz(pi) * roty(-pi/2)).T\n",
    "RguT = Rgu.T\n",
    "pprint(RguT)\n",
    "print(Rgu == RguT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#  euler_R is the composite rotation matrix of the following\n",
    "# a rotation of alpha in the z axis\n",
    "# a rotation of beta in the new y axis\n",
    "# a rotation of gamma in the new x axis \n",
    "# this will be useful later \n",
    "\n",
    "alpha, beta, gamma = symbols('alpha beta gamma', real = True)\n",
    "euler_R = rotz(alpha) * roty(beta) * rotx(gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_wrist_center(gripper_point, R0g, dg = 0.303):\n",
    "\n",
    "  # get the coordinates of the wrist center wrt to the base frame (xw, yw, zw)\n",
    "  # given the following info:\n",
    "  # the coordinates of the gripper (end effector) (x, y, z)\n",
    "  # the rotation of the gripper in gripper frame wrt to the base frame (R0u)\n",
    "  # the distance between gripper and wrist center dg which is along common z axis\n",
    "  xu, yu, zu = gripper_point \n",
    "    \n",
    "  nx, ny, nz = R0g[0, 2], R0g[1, 2], R0g[2, 2]\n",
    "  xw = xu - dg * nx\n",
    "  yw = yu - dg * ny\n",
    "  zw = zu - dg * nz \n",
    "\n",
    "  return xw, yw, zw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# This is given position and orientation of the gripper wrt to URDFrame\n",
    "px, py, pz = 0.49792, 1.3673, 2.4988\n",
    "roll, pitch, yaw = 0.366, -0.078, 2.561"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wrist_center (0.750499428337951, 1.20160389781975, 2.47518995758694)\n",
      "evaluated R0g:\n",
      "⎡0.257143295038827   0.48887208255965   -0.833595473062543⎤\n",
      "⎢                                                         ⎥\n",
      "⎢0.259329420712765  0.796053601157403    0.54685182237706 ⎥\n",
      "⎢                                                         ⎥\n",
      "⎣0.93092726749696   -0.356795110642117  0.0779209320563015⎦\n"
     ]
    }
   ],
   "source": [
    "gripper_point = px, py, pz\n",
    "\n",
    "# This is the rotation of the gripper in URDF wrt to base frame \n",
    "R0u_eval = euler_R.evalf(subs = {alpha: yaw, beta: pitch, gamma: roll})\n",
    "\n",
    "# R0g * Rgu = R0u \n",
    "R0g_eval = R0u_eval * RguT\n",
    "\n",
    "# calculate wrist center\n",
    "wrist_center = get_wrist_center(gripper_point, R0g_eval, dg = 0.303)\n",
    "print(\"wrist_center\", wrist_center)\n",
    "\n",
    "# evaluated R0g\n",
    "print(\"evaluated R0g:\")\n",
    "pprint(R0g_eval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_hypotenuse(a, b):\n",
    "  # calculate the longest side given the two shorter sides of a right triangle using pythagorean theorem\n",
    "  return sqrt(a*a + b*b)\n",
    "\n",
    "def get_cosine_law_angle(a, b, c):\n",
    "  # given all sides of a triangle a, b, c\n",
    "  # calculate angle gamma between sides a and b  using cosine law\n",
    "    \n",
    "  cos_gamma = (a*a + b*b - c*c) / (2*a*b)\n",
    "  sin_gamma = sqrt(1 - cos_gamma * cos_gamma)\n",
    "  gamma = atan2(sin_gamma, cos_gamma)\n",
    "\n",
    "  return gamma\n",
    "\n",
    "def get_first_three_angles(wrist_center):\n",
    "  # given the wrist center which a tuple of 3 numbers x, y, z\n",
    "  # (x, y, z) is the wrist center point wrt base frame\n",
    "  # return the angles q1, q2, q3 for each respective joint\n",
    "  # given geometry of the kuka kr210\n",
    "    \n",
    "  x, y, z  = wrist_center\n",
    "    \n",
    "  a1, a2, a3 = 0.35, 1.25, -0.054\n",
    "  d1, d4 = 0.75, 1.5\n",
    "  l = 1.50097168527591 #get_hypotenuse(d4, -a3)\n",
    "  phi = 1.53481186671284 # atan2(d4, -a3)\n",
    "  \n",
    "  x_prime = get_hypotenuse(x, y)\n",
    "  mx = x_prime -  a1\n",
    "  my = z - d1 \n",
    "  m = get_hypotenuse(mx, my)\n",
    "  alpha = atan2(my, mx)\n",
    "  \n",
    "  gamma = get_cosine_law_angle(l, a2, m)\n",
    "  beta = get_cosine_law_angle(m, a2, l)\n",
    "  \n",
    "  q1 = atan2(y, x)\n",
    "  q2 = pi/2 - beta - alpha \n",
    "  q3 = -(gamma - phi)\n",
    "    \n",
    "  return q1, q2, q3 \n",
    "    \n",
    "  return q1, q2, q3 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q1: 1.01249809363771\n",
      "q2: -0.275800363737724\n",
      "q3: -0.115686651053751\n"
     ]
    }
   ],
   "source": [
    "j1, j2, j3 = get_first_three_angles(wrist_center)\n",
    "\n",
    "print(\"q1:\", j1.evalf())\n",
    "print(\"q2:\", j2.evalf())\n",
    "print(\"q3:\", j3.evalf())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "'''\n",
    "Recall that from our simplification earlier, R36 equals the following:\n",
    "\n",
    "Matrix([\n",
    "[-sin(q4)*sin(q6) + cos(q4)*cos(q5)*cos(q6), -sin(q4)*cos(q6) - sin(q6)*cos(q4)*cos(q5), -sin(q5)*cos(q4)],\n",
    "[                           sin(q5)*cos(q6),                           -sin(q5)*sin(q6),          cos(q5)],\n",
    "[-sin(q4)*cos(q5)*cos(q6) - sin(q6)*cos(q4),  sin(q4)*sin(q6)*cos(q5) - cos(q4)*cos(q6),  sin(q4)*sin(q5)]])\n",
    "\n",
    "From trigonometry we can get q4, q5, q6 if we know numerical values of all cells of matrix R36  \n",
    "'''\n",
    "\n",
    "def get_last_three_angles(R):\n",
    "    \n",
    "  sin_q4 = R[2, 2]\n",
    "  cos_q4 =  -R[0, 2]\n",
    "    \n",
    "  sin_q5 = sqrt(R[0, 2]**2 + R[2, 2]**2) \n",
    "  cos_q5 = R[1, 2]\n",
    "    \n",
    "  sin_q6 = -R[1, 1]\n",
    "  cos_q6 = R[1, 0] \n",
    "  \n",
    "  q4 = atan2(sin_q4, cos_q4)\n",
    "  q5 = atan2(sin_q5, cos_q5)\n",
    "  q6 = atan2(sin_q6, cos_q6)\n",
    "    \n",
    "  return q4, q5, q6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    " - R0g = R03 * R36 * R6g\n",
    " - frame of joint 6 is the same orientation of gripper frame so  R6g = I\n",
    " - R03.T * R0g = R03.T * R03 * R36 * I\n",
    " ---> R36 = R03.T * R0g\n",
    " \n",
    " Recall we have this expression earlier for R03T:\n",
    "   Matrix([\n",
    "    [sin(q2 + q3)*cos(q1), sin(q1)*sin(q2 + q3),  cos(q2 + q3)],\n",
    "    [cos(q1)*cos(q2 + q3), sin(q1)*cos(q2 + q3), -sin(q2 + q3)],\n",
    "    [            -sin(q1),              cos(q1),             0]])\n",
    "\n",
    "\n",
    " Recall we also have evaluated R0g earlier.\n",
    "   Matrix([\n",
    "     [0.257143295038827, 0.488872082559650, -0.833595473062543],\n",
    "     [0.259329420712765, 0.796053601157403, 0.546851822377060], \n",
    "     [0.930927267496960, -0.356795110642117, 0.0779209320563015]])\n",
    "     \n",
    "\n",
    "  We also have solved for q1, q2, q3 earlier:\n",
    "    q1: 1.01249809363771\n",
    "    q2: -0.275800363737724\n",
    "    q3: -0.115686651053748\n",
    "\n",
    "  So we can actually evaluate for R36 because we have numerical values for\n",
    "    R03.T and R0g\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q1: 1.01249809363771\n",
      "q2: -0.275800363737724\n",
      "q3: -0.115686651053751\n",
      "q4: 1.63446527240323\n",
      "q5: 1.52050002599430\n",
      "q6: -0.815781306199679\n"
     ]
    }
   ],
   "source": [
    "R03T_eval = R03T.evalf(subs = {q1: j1.evalf(), q2: j2.evalf(), q3: j3.evalf()})\n",
    "R36_eval = R03T_eval * R0g_eval\n",
    "\n",
    "j4, j5, j6 = get_last_three_angles(R36_eval)\n",
    "\n",
    "print(\"q1:\", j1.evalf())\n",
    "print(\"q2:\", j2.evalf())\n",
    "print(\"q3:\", j3.evalf())\n",
    "print(\"q4:\", j4.evalf())\n",
    "print(\"q5:\", j5.evalf())\n",
    "print(\"q6:\", j6.evalf())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
